package com.tool.engine.config;

import brave.Span;
import brave.Tracer;
import brave.propagation.TraceContext;
import com.tool.engine.export.task.TaskDecorator;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
@ConditionalOnClass(name = "brave.Tracing")
public class ZipkinTraceConfiguration {

    @Bean
    @ConditionalOnClass(TaskDecorator.class)
    public TaskDecorator zipkinTraceDecorator(Tracer tracer) {
        return new ZipkinTraceContextDecorator(tracer);
    }

    static class ZipkinTraceContextDecorator implements TaskDecorator {
        private final Tracer tracer;

        public ZipkinTraceContextDecorator(Tracer tracer) {
            this.tracer = tracer;
        }

        @Override
        public Runnable decorate(Runnable runnable) {
            TraceContext context = tracer.currentSpan().context();
            return () -> {
                Tracer.SpanInScope spanInScope = null;
                try {
                    Span span = tracer.toSpan(context);
                    spanInScope = tracer.withSpanInScope(span);
                    runnable.run();
                }finally {
                    assert spanInScope != null;
                    spanInScope.close();
                }
            };
        }
    }
}
